import datetime
import os

import torch
import matplotlib
matplotlib.use('Agg')
import scipy.signal
from matplotlib import pyplot as plt
# from torch.utils.tensorboard import SummaryWriter

class LossHistory():
    def __init__(self, log_dir):
        time_str        = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')
        self.log_dir    = os.path.join(log_dir, "loss_psnr_" + str(time_str))
        self.psnr_train  = []
        self.psnr_val  = []
        self.losses     = []
        
        os.makedirs(self.log_dir,exist_ok=True)
        # self.writer     = SummaryWriter(self.log_dir)
        # dummy_input     = torch.randn(2, 3, input_shape[0], input_shape[1]).to('cuda')
        # self.writer.add_graph(model, dummy_input)

    def append_loss(self, psnr_t, loss, psnr_v):
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir, exist_ok=True)
            
        self.psnr_train.append(psnr_t)
        self.psnr_val.append(psnr_v)
        self.losses.append(loss)

        # with open(os.path.join(self.log_dir, "train_epoch_psnr.txt"), 'a') as f:
        #     f.write(str(psnr))
        #     f.write("\n")

        # with open(os.path.join(self.log_dir, "epoch_loss.txt"), 'a') as f:
        #     f.write(str(loss))
        #     f.write("\n")

        self.loss_plot()

    def loss_plot(self):
        iters = range(len(self.losses))

        plt.figure()
        # plt.plot(x, y, 'red', linewidth = 2, label='train loss')
        plt.plot(iters, self.losses, 'red', linewidth = 2, label='train loss')
        try:
            if len(self.losses) < 25:
                num = 5
            else:
                num = 15
            plt.plot(iters, scipy.signal.savgol_filter(self.losses, num, 3), 'green', linestyle = '--', linewidth = 2, label='smooth train loss')
        except:
            pass
        
        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend(loc="upper left")
        plt.savefig(os.path.join(self.log_dir, "epoch_loss.png"))
        plt.cla()
        plt.close("all")

        plt.figure()
        plt.plot(iters, self.psnr_train, 'red', linewidth = 2, label='train psnr')
        try:
            if len(self.psnr_train) < 25:
                num = 5
            else:
                num = 15
            plt.plot(iters, scipy.signal.savgol_filter(self.psnr_train, num, 3), 'green', linestyle = '--', linewidth = 2, label='smooth psnr train')
        except:
            pass
        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('PSNR train')
        plt.legend(loc="upper left")
        plt.savefig(os.path.join(self.log_dir, "epoch_pnsr_train.png"))
        plt.cla()
        plt.close("all")

        plt.figure()
        plt.plot(iters, self.psnr_val, 'red', linewidth = 2, label='val psnr')
        try:
            if len(self.psnr_val) < 25:
                num = 5
            else:
                num = 15
            plt.plot(iters, scipy.signal.savgol_filter(self.psnr_val, num, 3), 'green', linestyle = '--', linewidth = 2, label='smooth psnr val')
        except:
            pass
        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('PSNR val')
        plt.legend(loc="upper left")
        plt.savefig(os.path.join(self.log_dir, "epoch_pnsr_val.png"))
        plt.cla()
        plt.close("all")
